# -*- coding: utf-8 -*-

from fastapi import APIRouter, HTTPException, File, UploadFile, Form, Depends
from fastapi.responses import StreamingResponse
from typing import Dict, Any
import tempfile
import os
import json
import time
from ..services.chat_service import ChatService
from ..config import settings

router = APIRouter(tags=["chat"])

# 依赖注入：获取聊天服务实例
def get_chat_service() -> ChatService:
    """提供聊天服务实例的依赖注入函数"""
    return ChatService()

def get_file_extension(filename: str) -> str:
    """获取文件扩展名"""
    return os.path.splitext(filename)[1].lower()

@router.post("/upload-document", response_model=Dict[str, Any])
async def upload_document(
    session_id: str = Form(..., description="会话ID"),
    file: UploadFile = File(..., description="要上传的文件"),
    chat_service: ChatService = Depends(get_chat_service)
):
    """上传文档并添加到向量数据库"""
    try:
        # 保存文件到会话目录
        session_dir = os.path.join(settings.DOCUMENTS_DIR, session_id)
        os.makedirs(session_dir, exist_ok=True)
        
        file_path = os.path.join(session_dir, f"{os.urandom(8).hex()}{get_file_extension(file.filename)}")
        
        with open(file_path, "wb") as f:
            f.write(await file.read())
        
        # 上传文档
        result = chat_service.upload_document(
            session_id=session_id,
            file_path=file_path,
            file_name=file.filename,
            file_type=file.content_type
        )
        
        return result
    
    except Exception as e:
        # 记录详细错误信息以便调试
        error_msg = f"上传失败: {str(e)}"
        print(error_msg)
        raise HTTPException(status_code=500, detail=error_msg)

@router.post("/chat", response_class=StreamingResponse)
async def chat(
    session_id: str = Form(..., description="会话ID"),
    message: str = Form(..., description="用户消息"),
    chat_service: ChatService = Depends(get_chat_service)
):
    """发送聊天消息并获取流式响应"""
    try:
        # 定义流式响应生成器
        def generate_response():
            """生成流式响应的生成器函数"""
            # 初始化响应
            yield json.dumps({"status": "start", "content": ""}) + "\n"
            
            # 获取聊天响应
            response = chat_service.chat(session_id=session_id, message=message)
            answer = response.get("answer", "")
            
            # 流式输出
            chunk_size = 50  # 每次输出的字符数
            for i in range(0, len(answer), chunk_size):
                chunk = answer[i:i+chunk_size]
                yield json.dumps({"status": "stream", "content": chunk}) + "\n"
                time.sleep(0.05)  # 控制输出速度，模拟自然打字效果
            
            # 添加引用信息
            if "source_documents" in response and response["source_documents"]:
                sources = []
                for doc in response["source_documents"]:
                    if hasattr(doc.metadata, 'source'):
                        sources.append({
                            "title": doc.metadata.get('title', '文档片段'), 
                            "source": doc.metadata['source']
                        })
                
                if sources:
                    sources_text = "\n\n**参考资料**:\n" + "\n".join(
                        [f"- [{s['title']}]({s['source']})" for s in sources]
                    )
                    yield json.dumps({"status": "stream", "content": sources_text}) + "\n"
            
            # 结束响应
            yield json.dumps({"status": "end", "content": ""}) + "\n"
        
        return StreamingResponse(
            generate_response(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive"
            }
        )
    
    except Exception as e:
        error_message = f"聊天处理失败: {str(e)}"
        print(error_message)  # 记录错误
        error_response = json.dumps({"status": "error", "content": error_message}) + "\n"
        return StreamingResponse(
            [error_response], 
            media_type="text/event-stream"
        )    